// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include <cmath>
#include <krpc/rpc/errno.pb.h>
#include <krpc/rpc/config.h>
#include <krpc/rpc/policy/auto_concurrency_limiter.h>


namespace krpc {
    namespace policy {

        AutoConcurrencyLimiter::AutoConcurrencyLimiter()
                : _max_concurrency(turbo::get_flag(FLAGS_auto_cl_initial_max_concurrency)),
                  _remeasure_start_us(NextResetTime(kutil::gettimeofday_us())), _reset_latency_us(0),
                  _min_latency_us(-1), _ema_max_qps(-1),
                  _explore_ratio(turbo::get_flag(FLAGS_auto_cl_max_explore_ratio)), _last_sampling_time_us(0),
                  _total_succ_req(0) {
        }

        AutoConcurrencyLimiter *AutoConcurrencyLimiter::New(const AdaptiveMaxConcurrency &) const {
            return new(std::nothrow) AutoConcurrencyLimiter;
        }

        bool AutoConcurrencyLimiter::OnRequested(int current_concurrency, Controller *) {
            return current_concurrency <= _max_concurrency;
        }

        void AutoConcurrencyLimiter::OnResponded(int error_code, int64_t latency_us) {
            if (0 == error_code) {
                _total_succ_req.fetch_add(1, kutil::memory_order_relaxed);
            } else if (ELIMIT == error_code) {
                return;
            }

            const int64_t now_time_us = kutil::gettimeofday_us();
            int64_t last_sampling_time_us =
                    _last_sampling_time_us.load(kutil::memory_order_relaxed);

            if (last_sampling_time_us == 0 ||
                now_time_us - last_sampling_time_us >=
                turbo::get_flag(FLAGS_auto_cl_sampling_interval_ms) * 1000) {
                bool sample_this_call = _last_sampling_time_us.compare_exchange_strong(
                        last_sampling_time_us, now_time_us, kutil::memory_order_relaxed);
                if (sample_this_call) {
                    bool sample_window_submitted = AddSample(error_code, latency_us,
                                                             now_time_us);
                    if (sample_window_submitted) {
                        // The following log prints has data-race in extreme cases,
                        // unless you are in debug, you should not open it.
                        VLOG(1)
                                                << "Sample window submitted, current max_concurrency:"
                                                << _max_concurrency
                                                << ", min_latency_us:" << _min_latency_us
                                                << ", ema_max_qps:" << _ema_max_qps
                                                << ", explore_ratio:" << _explore_ratio;
                    }
                }
            }
        }

        int AutoConcurrencyLimiter::MaxConcurrency() {
            return _max_concurrency;
        }

        int64_t AutoConcurrencyLimiter::NextResetTime(int64_t sampling_time_us) {
            int64_t reset_start_us = sampling_time_us +
                                     (turbo::get_flag(FLAGS_auto_cl_noload_latency_remeasure_interval_ms) / 2 +
                                      kutil::fast_rand_less_than(
                                              turbo::get_flag(FLAGS_auto_cl_noload_latency_remeasure_interval_ms) /
                                              2)) * 1000;
            return reset_start_us;
        }

        bool AutoConcurrencyLimiter::AddSample(int error_code,
                                               int64_t latency_us,
                                               int64_t sampling_time_us) {
            std::unique_lock<kutil::Mutex> lock_guard(_sw_mutex);
            if (_reset_latency_us != 0) {
                // min_latency is about to be reset soon.
                if (_reset_latency_us > sampling_time_us) {
                    // ignoring samples during waiting for the deadline.
                    return false;
                }
                // Remeasure min_latency when concurrency has dropped to low load
                _min_latency_us = -1;
                _reset_latency_us = 0;
                _remeasure_start_us = NextResetTime(sampling_time_us);
                ResetSampleWindow(sampling_time_us);
            }

            if (_sw.start_time_us == 0) {
                _sw.start_time_us = sampling_time_us;
            }

            if (error_code != 0 && turbo::get_flag(FLAGS_auto_cl_enable_error_punish)) {
                ++_sw.failed_count;
                _sw.total_failed_us += latency_us;
            } else if (error_code == 0) {
                ++_sw.succ_count;
                _sw.total_succ_us += latency_us;
            }

            if (_sw.succ_count + _sw.failed_count < turbo::get_flag(FLAGS_auto_cl_min_sample_count)) {
                if (sampling_time_us - _sw.start_time_us >=
                    turbo::get_flag(FLAGS_auto_cl_sample_window_size_ms) * 1000) {
                    // If the sample size is insufficient at the end of the sampling
                    // window, discard the entire sampling window
                    ResetSampleWindow(sampling_time_us);
                }
                return false;
            }
            if (sampling_time_us - _sw.start_time_us <
                turbo::get_flag(FLAGS_auto_cl_sample_window_size_ms) * 1000 &&
                _sw.succ_count + _sw.failed_count < turbo::get_flag(FLAGS_auto_cl_max_sample_count)) {
                return false;
            }

            if (_sw.succ_count > 0) {
                UpdateMaxConcurrency(sampling_time_us);
            } else {
                // All request failed
                AdjustMaxConcurrency(_max_concurrency / 2);
            }
            ResetSampleWindow(sampling_time_us);
            return true;
        }

        void AutoConcurrencyLimiter::ResetSampleWindow(int64_t sampling_time_us) {
            _total_succ_req.exchange(0, kutil::memory_order_relaxed);
            _sw.start_time_us = sampling_time_us;
            _sw.succ_count = 0;
            _sw.failed_count = 0;
            _sw.total_failed_us = 0;
            _sw.total_succ_us = 0;
        }

        void AutoConcurrencyLimiter::UpdateMinLatency(int64_t latency_us) {
            const double ema_factor = turbo::get_flag(FLAGS_auto_cl_alpha_factor_for_ema);
            if (_min_latency_us <= 0) {
                _min_latency_us = latency_us;
            } else if (latency_us < _min_latency_us) {
                _min_latency_us = latency_us * ema_factor + _min_latency_us * (1 - ema_factor);
            }
        }

        void AutoConcurrencyLimiter::UpdateQps(double qps) {
            const double ema_factor = turbo::get_flag(FLAGS_auto_cl_alpha_factor_for_ema) / 10;
            if (qps >= _ema_max_qps) {
                _ema_max_qps = qps;
            } else {
                _ema_max_qps = qps * ema_factor + _ema_max_qps * (1 - ema_factor);
            }
        }

        void AutoConcurrencyLimiter::AdjustMaxConcurrency(int next_max_concurrency) {
            next_max_concurrency = std::max(turbo::get_flag(FLAGS_kthread_concurrency), next_max_concurrency);
            if (next_max_concurrency != _max_concurrency) {
                _max_concurrency = next_max_concurrency;
            }
        }

        void AutoConcurrencyLimiter::UpdateMaxConcurrency(int64_t sampling_time_us) {
            int32_t total_succ_req = _total_succ_req.load(kutil::memory_order_relaxed);
            double failed_punish = _sw.total_failed_us * turbo::get_flag(FLAGS_auto_cl_fail_punish_ratio);
            int64_t avg_latency =
                    std::ceil((failed_punish + _sw.total_succ_us) / _sw.succ_count);
            double qps = 1000000.0 * total_succ_req / (sampling_time_us - _sw.start_time_us);
            UpdateMinLatency(avg_latency);
            UpdateQps(qps);

            int next_max_concurrency = 0;
            // Remeasure min_latency at regular intervals
            if (_remeasure_start_us <= sampling_time_us) {
                const double reduce_ratio = turbo::get_flag(FLAGS_auto_cl_reduce_ratio_while_remeasure);
                _reset_latency_us = sampling_time_us + avg_latency * 2;
                next_max_concurrency =
                        std::ceil(_ema_max_qps * _min_latency_us / 1000000 * reduce_ratio);
            } else {
                const double change_step = turbo::get_flag(FLAGS_auto_cl_change_rate_of_explore_ratio);
                const double max_explore_ratio = turbo::get_flag(FLAGS_auto_cl_max_explore_ratio);
                const double min_explore_ratio = turbo::get_flag(FLAGS_auto_cl_min_explore_ratio);
                const double correction_factor = turbo::get_flag(FLAGS_auto_cl_latency_fluctuation_correction_factor);
                if (avg_latency <= _min_latency_us * (1.0 + min_explore_ratio * correction_factor) ||
                    qps <= _ema_max_qps / (1.0 + min_explore_ratio)) {
                    _explore_ratio = std::min(max_explore_ratio, _explore_ratio + change_step);
                } else {
                    _explore_ratio = std::max(min_explore_ratio, _explore_ratio - change_step);
                }
                next_max_concurrency =
                        _min_latency_us * _ema_max_qps / 1000000 * (1 + _explore_ratio);
            }

            AdjustMaxConcurrency(next_max_concurrency);
        }

    }  // namespace policy
}  // namespace krpc
